package b2_AVL;

import java.util.ArrayList;
import java.util.List;
import java.util.NoSuchElementException;

/**
 * 二叉平衡树
 */
public class AVLTree {
    private class TreeNode {
        int val;
        // 平衡因子 = 右子树高度 - 左子树的高度
        int bf;
        TreeNode left;
        TreeNode right;
        TreeNode parent;

        TreeNode(int val) {
            this.val = val;
        }
    }

    private TreeNode root;

    /**
     * 插入操作
     *   1. 先把数据插入二叉树中
     *   2. 根据平衡因子对树进行平衡调整
     * @return 无法插入重复元素
     */
    public boolean insert(int val) {
        TreeNode node = new TreeNode(val);
        if(root == null) {
            root = node;
            return true;
        }

        TreeNode parent = null;
        TreeNode cur = root;
        while(cur != null) {
            if(val == cur.val){
                return false;
            }else if(val < cur.val) {
                parent = cur;
                cur = parent.left;
            }else{
                parent = cur;
                cur = parent.right;
            }
        }
        // 此时cur=null，插入 val
        if(val < parent.val) {
            parent.left = node;
        }else {
            parent.right = node;
        }
        node.parent = parent;
        cur = node;

        /**
         * 自底向上调节平衡因子
         *   因为 bf = rightH - leftH, 所以当某个子树插入一个节点不引起左右子树高度变化时，就不需要继续向上调整平衡因子了
         *   也就是，当某个子树节点从 (-1 / 1) -> 0 的时候
         */
        while(parent != null) {
            // 一直向上调节
            if(cur == parent.left) {
                // cur 是 parent 的左孩子
                parent.bf --;
            }else {
                parent.bf ++;
            }

            if(parent.bf == 0) {
                // 已经平衡了
                break;
            }else if(parent.bf == -1 || parent.bf == 1) {
                // 继续向上
                cur = parent;
                parent = cur.parent;
            }else {
                // parent.bf = -2/2, 该子树不平衡了，应该旋转, 有以下4种情况：
                if(parent.bf == -2 && cur.bf == -1) {
                    // LL型，右单旋
                    rotateR(parent);
                }else if(parent.bf == 2 && cur.bf == 1) {
                    // RR型，左单旋
                    rotateL(parent);
                }else if(parent.bf == -2 && cur.bf == 1) {
                    // LR型，左右双旋
                    rotateLR(parent);
                }else {
                    // parent.bf == 2 && cur.bf == -1
                    // RL型，右左双旋
                    rotateRL(parent);
                }
                // 旋转完就一定平衡了
                break;
            }
        }
        return true;
    }
    // 左单旋
    private void rotateL(TreeNode p) {
        // p 的母节点
        TreeNode pp = p.parent;
        // p 的右孩子
        TreeNode subR = p.right;
        // subR 的左孩子，可能不存在
        TreeNode subRL = subR.left;

        // subR 提上去
        if(pp == null) {
            this.root = subR;
        }else if(pp.left == p) {
            pp.left = subR;
        }else {
            // pp.right == parent
            pp.right = subR;
        }
        subR.parent = pp;
        // p 作为 subR 的左孩子
        subR.left = p;
        p.parent = subR;
        // p 与 subRL 连接
        p.right = subRL;
        if(subRL != null) {
            subRL.parent = p;
        }

        // 经过旋转最终只会改变 parent 和 subR 的平衡因子
        // 而且无论 subRL 是否存在，最终都是 0
        subR.bf = 0;
        p.bf = 0;
    }
    // 右单旋
    private void rotateR(TreeNode p) {
        // p 的父节点
        TreeNode pp = p.parent;
        // p 的左孩子
        TreeNode subL = p.left;
        // subL 的右孩子，可能不存在
        TreeNode subLR = subL.right;

        if(pp == null) {
            this.root = subL;
        }else if(pp.left == p) {
            pp.left = subL;
        }else {
            pp.right = subL;
        }
        subL.parent = pp;
        subL.right = p;
        p.parent = subL;
        p.left = subLR;
        if(subLR != null) {
            subLR.parent = p;
        }
        // 经过旋转最终只会改变 parent 和 subL 的平衡因子
        // 而且无论 subLR 是否存在，最终都是 0
        subL.bf = 0;
        p.bf = 0;
    }
    // 左右双旋
    private void rotateLR(TreeNode p) {
        TreeNode subL = p.left;
        TreeNode subLR = subL.right;
        int bf = subLR.bf; // 可能为 0、1、-1

        rotateL(subL);
        rotateR(p);

        subLR.bf = 0;
        if(bf == 1) {
            p.bf = 0;
            subL.bf = -1;
        }else if(bf == -1) {
            subL.bf = 0;
            p.bf = 1;
        }
        // 当 bf == 0 时，就不需要再手动改平衡因子了
    }
    // 右左双旋
    private void rotateRL(TreeNode p) {
        TreeNode subR = p.right;
        TreeNode subRL = subR.left;
        int bf = subRL.bf;

        rotateR(subR);
        rotateL(p);

        subRL.bf = 0;
        if(bf == -1) {
            p.bf = 0;
            subR.bf = 1;
        }else if(bf == 1) {
            subR.bf = 0;
            p.bf = -1;
        }
    }

    /**
     * 查找操作
     */
    public boolean contains(int val) {
        return contains(root, val);
    }
    // 在当前树中查找元素
    private boolean contains(TreeNode root, int val) {
        if(root == null) {
            return false;
        }
        if(val == root.val) {
            return true;
        }else if(val < root.val) {
            return contains(root.left, val);
        }else {
            return contains(root.right, val);
        }
    }

    /**
     * 删除操作
     * Hibbard Deletion: 拿前驱 / 后继 来代替
     */
    public boolean delete(int val) {
        if(root == null) {
            return false;
        }
        // ......
        return true;
    }

    /** 打印二叉树, 中序遍历的结果 **/
    @Override
    public String toString() {
        List<Integer> list = new ArrayList<>();
        inOrder(list, root);
        return list.toString() + "\n" + "是否平衡: " + isBalance(this.root);
    }
    // 中序遍历
    private void inOrder(List<Integer> list, TreeNode root) {
        if(root == null){
            return;
        }
        inOrder(list, root.left);
        list.add(root.val);
        inOrder(list, root.right);
    }

    /**
     * 检验二叉树是否平衡
     *   每个节点子树高度差的绝对值不超过1(注意节点中如果没有平衡因子)
     *   节点的平衡因子是否计算正确
     */
    private boolean isBalance(TreeNode root) {
        if(root == null) {
            return true;
        }
        int leftH = height(root.left);
        int rightH = height(root.right);
        if(rightH - leftH != root.bf) {
            System.err.println("平衡因子异常！");
            return false;
        }
        return Math.abs(root.bf) < 2 && isBalance(root.left) && isBalance(root.right);
    }
    // 计算二叉树的高度
    private int height(TreeNode root) {
        if(root == null) {
            return 0;
        }
        return 1 + Math.max(height(root.left), height(root.right));
    }
}
